package leetcode.p930;

import java.util.ArrayList;
import java.util.List;

/**
 * @author: trtan
 * @date: 2021-07-09 17:00
 **/
public class BinarySubArraysWithSum {
    /**
     * 以1分段，分别记录每段0的数量
     * 求组成goal的方案数其实就是选择连续goal个1
     * 然后左右端点的1左右两边分别会存在一定数量的0
     * 那么这段方案数就是两边0数量的乘积
     * 于是可分情况求得答案
     * @param nums
     * @param goal
     * @return int
     * @author trtan
     * @date 2021/7/9 17:39
     */
    public int numSubarraysWithSum(int[] nums, int goal) {
        // list记录以1分割每段0的数量
        List<Integer> list = new ArrayList<>();
        int i = 0, now = 0, count_one = 0;
        while (i < nums.length) {
            if (now == 0) {
                int cnt = 0;
                while (i < nums.length && nums[i] == now) {
                    i++;
                    cnt++;
                }
                list.add(cnt);
            } else {
                i++;
                count_one++;
            }
            now = 1 - now;
        }
        // 确保最后一个是1时结尾0的数量为0
        if (now == 0) {
            list.add(0);
        }
        int ans = 0;
        //如果没有那么多数量的1那么方案数为0
        if (goal > count_one) {
            ans = 0;
        }
        else if (goal == 0) {
            // 如果目标是凑成0，那么特殊处理，求每段0能凑成的区间数量
            for (int j = 0; j < list.size(); j++) {
                int x = list.get(j);
                ans += (x * (x + 1) / 2);
            }
        } else {
            for (int j = goal; j <= count_one; j++) {
                int a = list.get(j - goal) + 1;
                int b = list.get(j) + 1;
                ans += a * b;
            }
        }
        return ans;
    }

    public static void main(String[] args) {
        System.out.println(new BinarySubArraysWithSum().numSubarraysWithSum(new int[]{1, 0, 1, 0, 1}, 2));
    }
}
